import numpy as np # type: ignore
import pandas as pd # type: ignore

def calc_p_diff_adjust_factors(tunnels = '', case = ''):
    raw_p_above_box_df = pd.read_csv('comsol_raw/p_above_box' + tunnels + '.csv', skiprows = 8, index_col = [0,1,2])
    raw_p_below_box_df = pd.read_csv('comsol_raw/p_below_box' + tunnels + '.csv', skiprows = 8, index_col = [0,1,2])
    raw_p_above_slice_df = pd.read_csv('comsol_raw/p_above_slice' + tunnels + case + '.csv', skiprows = 8, index_col = [0,1,2])
    raw_p_below_slice_df = pd.read_csv('comsol_raw/p_below_slice' + tunnels + case + '.csv', skiprows = 8, index_col = [0,1,2])

    freq_strs = list(raw_p_above_slice_df.columns.get_level_values(0))
    freqs = sorted(set([float(i.split('freq=')[1]) for i in freq_strs]))    

    new_box_df = pd.DataFrame(index = freqs, columns = ['p_sv_re', 'p_sv_im','p_st_re','p_st_im', 'p_diff_mag', 'p_diff_phase'])
    new_slice_df = pd.DataFrame(index = freqs, columns = ['p_sv_re', 'p_sv_im','p_st_re','p_st_im', 'p_diff_mag', 'p_diff_phase'])
    new_box_df.index.name = "freq"
    new_slice_df.index.name = "freq"
    # pt_box_above = raw_p_above_box_df.index[0]
    # pt_box_below = raw_p_below_box_df.index[0]
    # pt_slice_above = raw_p_above_slice_df.index[0]
    # pt_slice_below = raw_p_below_slice_df.index[0]


    for freqi in freqs:
        freqi = int(freqi)

        for new_df, raw_above_df, raw_below_df in zip([new_box_df, new_slice_df], [raw_p_above_box_df, raw_p_above_slice_df], [raw_p_below_box_df, raw_p_below_slice_df]):
    
            new_df.loc[freqi, 'p_sv_re'] = raw_above_df.loc[:, 'real(ta.p_t) (Pa) @ freq='+str(freqi)].iloc[0]
            new_df.loc[freqi, 'p_sv_im'] = raw_above_df.loc[:, 'imag(ta.p_t) (Pa) @ freq='+str(freqi)].iloc[0]
            
            new_df.loc[freqi, 'p_st_re'] = raw_below_df.loc[:, 'real(ta.p_t) (Pa) @ freq='+str(freqi)].iloc[0]
            new_df.loc[freqi, 'p_st_im'] = raw_below_df.loc[:, 'imag(ta.p_t) (Pa) @ freq='+str(freqi)].iloc[0]
        
    for new_df, model in zip( [new_box_df, new_slice_df], ['box', 'slice'] ):
        p_re_sv = np.asarray(new_df.loc[:, 'p_sv_re'])
        p_im_sv = np.asarray(new_df.loc[:, 'p_sv_im'])
        p_re_st = np.asarray(new_df.loc[:, 'p_st_re'])
        p_im_st = np.asarray(new_df.loc[:, 'p_st_im'])

        p_diff = np.asarray(p_re_sv - p_re_st + 1j*(p_im_sv - p_im_st))
        p_diff_mag = np.abs(p_diff)
        p_diff_phase = np.asarray([np.angle(p_diff[i]) for i in range(len(freqs))])

        new_df.loc[:, 'p_diff_mag'] = p_diff_mag
        new_df.loc[:, 'p_diff_phase'] = p_diff_phase

        new_df.to_csv('csv/p_diff_'+model+'.csv')

    p_diff_box = new_box_df.loc[:, 'p_diff_mag']
    p_diff_slice = new_slice_df.loc[:, 'p_diff_mag']

    adjust_factors = p_diff_box/p_diff_slice
    adjust_factors_df = pd.DataFrame(index = freqs, columns = ['adjust_factors'])
    adjust_factors_df.loc[:, 'adjust_factors'] = adjust_factors
    adjust_factors_df.to_csv('csv/p_adjust_factors' + tunnels + case + '.csv', header = False)

calc_p_diff_adjust_factors('', '_k_div_3')